import os

from langchain.chat_models import QianfanChatEndpoint
from langchain_core.language_models.chat_models import HumanMessage

os.environ["QIANFAN_AK"] = "aYSxDro6UlsB2ooAgNgAGBbK"
os.environ["QIANFAN_SK"] = "xoGnZ5GjrN5tUd2tm6T2hX7ofMc4LqXG"

chat = QianfanChatEndpoint(
    streaming=True,
)

# res = chat([HumanMessage(content="write a funny joke")])
# print(res)

res = chat.stream([HumanMessage(content="hi")], streaming=True)
for r in res:
    print(r)
    # print("chat resp:", r)